import json
import os
import sys
import re


def load_input(path):
    _, ext = os.path.splitext(path.lower())
    if ext in (".jsonl", ".jsonl"):
        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    yield json.loads(line)
    else:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
            if isinstance(data, list):
                for item in data:
                    yield item
            else:
                raise ValueError("Expected a list for JSON input")


def convert_answer(text):
    """
    Given text like:
      <think>…</think>
      <answer>…final answer…</answer>
    returns:
      <think>…</think>

      …final answer…
    """
    pattern = r"(<think>.*?</think>)\s*<answer>\s*(.*?)\s*</answer>"
    replacement = r"\1\n\n\2"
    return re.sub(pattern, replacement, text, flags=re.DOTALL)


def transform_record(rec):
    # print(rec)
    img = item['image']
    new_item = {
        "id": f"item_{img}",
        "image": img,
        "enable_thinking": True,
        "conversations": [
        {
        "from": "human",
        "value": item["instruction"]
        },
        {
        "from": "gpt",
        "value": f"<think>\n{item['thinking']}\n</think>\n\n{item['output']}"
        }
            ]
    }
    return new_item


def main():
    if len(sys.argv) != 2:
        print("Usage: python convert_to_qwen3_format.py <input_json_or_jsonl>")
        sys.exit(1)

    input_path = sys.argv[1]
    base, _ = os.path.splitext(input_path)
    output_path = f"{base}_qwen3_format.json"

    records = list(load_input(input_path))
    transformed = [transform_record(r) for r in records]

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(transformed, f, indent=2, ensure_ascii=False)

    print(f"Saved {len(transformed)} records to {output_path}")


if __name__ == "__main__":
    main()
